import pandas as pd


class CustomDataFrame(pd.DataFrame):
    # new(): read_csv() -> into_code_dict()
    @classmethod
    def new(cls, csv_path: str) -> dict:
        df = pd.read_csv(csv_path)
        return cls(df).into_code_dict()

    @classmethod
    def use_main(cls, csv_path: str) -> dict:
        df = pd.read_csv(csv_path)
        return cls(df).into_code_dict_main()

    @classmethod
    def into_code_list(cls, csv_path: str) -> list:
        df = pd.read_csv(csv_path)

        df["stock_code"] = df["stock_code"].astype(str).str.pad(6, "left", "0")
        return df["stock_code"].tolist()

    # pd.DataFrame.to_dict(): key: stock_code; value: item
    def into_code_dict(self) -> dict:
        self["stock_code"] = self["stock_code"].astype(str).str.pad(6, "left", "0")
        return {row["stock_code"]: row.to_dict() for _, row in self.iterrows()}

    def into_code_dict_main(self) -> dict:
        self["stock_code"] = self["stock_code"].astype(str).str.pad(6, "left", "0")
        filtered_df = self[self["stock_code"].str.startswith(("60", "000"))]
        return {row["stock_code"]: row.to_dict() for _, row in filtered_df.iterrows()}
